#include "common.cuh"
#include "ops.cuh"

void initalData(float *data, int size) {
    for (int i = 0; i < size; ++i) {
        data[i] = (float)(rand() & 0xFF) / 10.0f;
    }
}

__global__ void matrixSum_1D(const float *A, const float *B, float *C, int N) {
    int bid = blockIdx.x;
    int tid = threadIdx.x;
    int index = tid + bid * blockDim.x;

    C[index] = A[index] + B[index];
}

void matrixSum_1D_CPU(const float *A, const float *B, float *C, int N) {
    for (int i = 0; i < N; ++i) {
        C[i] = A[i] + B[i];
    }
}

int testSum_1D() {
    int iElementCount = 8192 * 32;
    size_t iSize = iElementCount * sizeof(float);

    float *fpHost_A, *fpHost_B, *fpHost_C_CPU, *fpHost_C_GPU;
    fpHost_A = (float *)malloc(iSize);
    fpHost_B = (float *)malloc(iSize);
    fpHost_C_CPU = (float *)malloc(iSize);
    fpHost_C_GPU = (float *)malloc(iSize);

    if (fpHost_A != nullptr && fpHost_B != nullptr && fpHost_C_CPU != nullptr &&
        fpHost_C_GPU != nullptr) {
        memset(fpHost_A, 0, iSize);
        memset(fpHost_B, 0, iSize);
        memset(fpHost_C_GPU, 0, iSize);
        memset(fpHost_C_CPU, 0, iSize);
    } else {
        printf("Failed to allocate host vectors!\n");
        free(fpHost_A);
        free(fpHost_B);
        free(fpHost_C_CPU);
        free(fpHost_C_GPU);
        return -1;
    }

    // 分配显存
    float *fpDevice_A, *fpDevice_B, *fpDevice_C;
    errorCheck(cudaMalloc((float **)&fpDevice_A, iSize), __FILE__, __LINE__);
    errorCheck(cudaMalloc((float **)&fpDevice_B, iSize), __FILE__, __LINE__);
    errorCheck(cudaMalloc((float **)&fpDevice_C, iSize), __FILE__, __LINE__);
    if (fpDevice_A != nullptr && fpDevice_B != nullptr &&
        fpDevice_C != nullptr) {
        cudaMemset(fpDevice_A, 0, iSize);
        cudaMemset(fpDevice_B, 0, iSize);
        cudaMemset(fpDevice_C, 0, iSize);
    } else {
        printf("Failed to allocate device vectors!\n");
        cudaFree(fpDevice_A);
        cudaFree(fpDevice_B);
        cudaFree(fpDevice_C);
        return -1;
    }

    // 初始化主机数据
    srand(42);
    initalData(fpHost_A, iElementCount);
    initalData(fpHost_B, iElementCount);

    // 复制数据到显存
    errorCheck(cudaMemcpy(fpDevice_A, fpHost_A, iSize, cudaMemcpyHostToDevice),
               __FILE__, __LINE__);
    errorCheck(cudaMemcpy(fpDevice_B, fpHost_B, iSize, cudaMemcpyHostToDevice),
               __FILE__, __LINE__);
    errorCheck(
        cudaMemcpy(fpDevice_C, fpHost_C_GPU, iSize, cudaMemcpyHostToDevice),
        __FILE__, __LINE__);

    // 在GPU中计算
    dim3 block(1024);
    dim3 grid(iElementCount / 1024);

    clock_t start, end;
    start = clock();
    matrixSum_1D<<<grid, block>>>(fpDevice_A, fpDevice_B, fpDevice_C,
                                  iElementCount);
    errorCheck(cudaGetLastError(), __FILE__, __LINE__);
    errorCheck(cudaDeviceSynchronize(), __FILE__, __LINE__);
    end = clock();

    // 复制结果回主机
    errorCheck(
        cudaMemcpy(fpHost_C_GPU, fpDevice_C, iSize, cudaMemcpyDeviceToHost),
        __FILE__, __LINE__);

    printf("Matrix sum 1D GPU Time elapsed: %f ms\n",
           (double)(end - start) * 1000.0 / CLOCKS_PER_SEC);

    start = clock();
    matrixSum_1D_CPU(fpHost_A, fpHost_B, fpHost_C_CPU, iElementCount);
    end = clock();

    printf("Matrix sum 1D CPU Time elapsed: %f ms\n",
           (double)(end - start) * 1000.0 / CLOCKS_PER_SEC);

    for (int i = 0; i < iElementCount; i++) {
        if (fabs(fpHost_C_CPU[i] - fpHost_C_GPU[i]) > 1e-5) {
            printf("Result verification failed at element %d!\n", i);
            return -1;
        }
    }

    free(fpHost_A);
    free(fpHost_B);
    free(fpHost_C_CPU);
    free(fpHost_C_GPU);

    errorCheck(cudaFree(fpDevice_A), __FILE__, __LINE__);
    errorCheck(cudaFree(fpDevice_B), __FILE__, __LINE__);
    errorCheck(cudaFree(fpDevice_C), __FILE__, __LINE__);

    errorCheck(cudaDeviceReset(), __FILE__, __LINE__);

    return 0;
}